#include <WiFi.h>
#include "codec.h"

static const char* wsHost     = "106.53.186.158";
static const uint16_t wsPort  = 8800;

WiFiClient wsClient;
static int wsConnected = 0;

static unsigned long lastConnectAttempt = 0;
static int connectRetryCount = 0;

void ws_network_diagnostics();

int ws_connect() {
  wsConnected = 0;
  
  if (wsClient.connected()) {
    wsClient.stop();
  }
  
  unsigned long RECONNECT_INTERVAL;
  if (connectRetryCount < 5) {
    RECONNECT_INTERVAL = 2000;
  } else if (connectRetryCount < 10) {
    RECONNECT_INTERVAL = 5000;
  } else {
    RECONNECT_INTERVAL = 10000;
  }
  
  unsigned long now = millis();
  if (now - lastConnectAttempt < RECONNECT_INTERVAL) {
    return -1;
  }
  
  lastConnectAttempt = now;
  connectRetryCount++;
  
  Serial.printf("connect to server %s:%d (%d)...", wsHost, wsPort, connectRetryCount);
  
  if (!wsClient.connect(wsHost, wsPort)) {
    Serial.println("connect failed, retry...");
    
    if (connectRetryCount == 3) {
      ws_network_diagnostics();
    }
    
    return -1;
  }
  
  wsClient.setNoDelay(true);

  Serial.println("connect success!");
  connectRetryCount = 0;

  uint8_t key[16];
  for (int i = 0; i < 16; i++) {
    key[i] = random(256);
  }
  
  const char* base64_chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  char wsKey[25];
  int j = 0;
  for (int i = 0; i < 16; i += 3) {
    uint32_t val = (key[i] << 16) | ((i+1 < 16 ? key[i+1] : 0) << 8) | (i+2 < 16 ? key[i+2] : 0);
    wsKey[j++] = base64_chars[(val >> 18) & 0x3F];
    wsKey[j++] = base64_chars[(val >> 12) & 0x3F];
    wsKey[j++] = base64_chars[(val >> 6) & 0x3F];
    wsKey[j++] = base64_chars[val & 0x3F];
  }
  wsKey[22] = '=';
  wsKey[23] = '=';
  wsKey[24] = '\0';

  wsClient.print("GET "); wsClient.print("/"); wsClient.println(" HTTP/1.1");
  wsClient.print("Host: "); wsClient.println(wsHost);
  wsClient.println("Upgrade: websocket");
  wsClient.println("Connection: Upgrade");
  wsClient.print("Sec-WebSocket-Key: "); wsClient.println(wsKey);
  wsClient.println("Sec-WebSocket-Version: 13");
  
  uint8_t mac[6];
  WiFi.macAddress(mac);
  char deviceId[32];
  char clientId[32];
  sprintf(deviceId, "ESP32C3_%02X%02X%02X", mac[3], mac[4], mac[5]);
  sprintf(clientId, "OPUS_%02X%02X%02X%02X%02X%02X", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  
  IPAddress localIP = WiFi.localIP();
  
  wsClient.print("Device-Id: "); wsClient.println(deviceId);
  wsClient.print("Client-Id: "); wsClient.println(clientId);
  wsClient.print("Local-IP: "); wsClient.println(localIP.toString());
  wsClient.println("Channels: 1");

  char sampleRate[32];
  sprintf(sampleRate, "Sample-Rate: %d", I2S_SAMPLE_RATE);
  wsClient.println(sampleRate);
  char bitrateStr[32];
  sprintf(bitrateStr, "Bitrate: %d", AUDIO_OPUS_BITRATE);
  wsClient.println(bitrateStr);


  char frameDuration[32];
  sprintf(frameDuration, "Frame-Duration: %d", AUDIO_OPUS_FRAME_MS);
  wsClient.println(frameDuration);

  // wsClient.println("Segment-Frames: 90");
  wsClient.println("Bit-Depth: 16");
  wsClient.println("Origin: ESP32C3");
  
  wsClient.println();

  String line = wsClient.readStringUntil('\n');
  if (line.indexOf("101") == -1) {
    Serial.println("handshake failed");
    wsClient.stop();
    return -2;
  }

  while (wsClient.connected()) {
    line = wsClient.readStringUntil('\n');
    if (line == "\r" || line.length() == 0) break;
  }

  wsConnected = 1;
  Serial.println("WebSocket connect success");
  return 0;
}

bool ws_is_connected() {
  return wsConnected && wsClient.connected();
}

void ws_reset_connection() {
  wsConnected = 0;
  if (wsClient.connected()) {
    wsClient.stop();
  }
  connectRetryCount = 0;
  lastConnectAttempt = 0;
  Serial.println("WebSocket reset");
}

int ws_get_retry_count() {
  return connectRetryCount;
}

static unsigned long lastHealthCheck = 0;
static int sendFailureCount = 0;

void ws_check_connection_health() {
  unsigned long now = millis();
  if (now - lastHealthCheck > 5000) {
    if (wsConnected && !wsClient.connected()) {
      Serial.println("connection health check failed, reset");
      wsConnected = 0;
      sendFailureCount = 0;
    }
    lastHealthCheck = now;
  }
}

void ws_network_diagnostics() {
  Serial.println("=== Network Diagnostics ===");
  
  // WiFi状态
  if (WiFi.status() == WL_CONNECTED) {
    Serial.printf("WiFi: Connected to %s\n", WiFi.SSID().c_str());
    Serial.printf("Local IP: %s\n", WiFi.localIP().toString().c_str());
    Serial.printf("Gateway: %s\n", WiFi.gatewayIP().toString().c_str());
    Serial.printf("DNS: %s\n", WiFi.dnsIP().toString().c_str());
    Serial.printf("RSSI: %d dBm\n", WiFi.RSSI());
  } else {
    Serial.printf("WiFi: Disconnected (status: %d)\n", WiFi.status());
    return;
  }
  
  // 测试DNS解析
  Serial.printf("DNS resolve %s... ", wsHost);
  IPAddress serverIP;
  if (WiFi.hostByName(wsHost, serverIP)) {
    Serial.printf("OK -> %s\n", serverIP.toString().c_str());
    
    // 尝试连接测试
    Serial.printf("Connection test to %s:%d... ", wsHost, wsPort);
    WiFiClient testClient;
    if (testClient.connect(wsHost, wsPort)) {
      Serial.println("OK");
      testClient.stop();
    } else {
      Serial.println("FAILED - server unreachable");
    }
  } else {
    Serial.println("FAILED - DNS resolution failed");
  }
  
  Serial.println("=========================");
}

void ws_send_bytes(uint8_t* buff, int length) {
  if (!ws_is_connected()) {
    ws_connect();
    if (!ws_is_connected()) {
      return;
    }
  }
  
  // 发送前再次检查连接状态
  if (!wsClient.connected()) {
    Serial.println("connection lost before sending");
    wsConnected = 0;
    return;
  }

  uint8_t header[14];
  int headerSize = 0;

  header[headerSize++] = 0x82; // FIN=1, 二进制

  // 客户端必须设置MASK位
  if (length <= 125) {
    header[headerSize++] = 0x80 | length; // MASK=1, payload长度
  } else if (length <= 65535) {
    header[headerSize++] = 0x80 | 126; // MASK=1, 扩展长度
    header[headerSize++] = (length >> 8) & 0xFF;
    header[headerSize++] = length & 0xFF;
  } else {
    header[headerSize++] = 0x80 | 127; // MASK=1, 64位长度
    for (int i = 7; i >= 0; i--) {
      header[headerSize++] = (length >> (8 * i)) & 0xFF;
    }
  }

  // 生成4字节随机mask key
  uint8_t maskKey[4];
  for (int i = 0; i < 4; i++) {
    maskKey[i] = random(256);
    header[headerSize++] = maskKey[i];
  }

  // 发送头部
  if (wsClient.write(header, headerSize) != headerSize) {
    Serial.println("send header failed, reset connection");
    wsConnected = 0;
    return;
  }
  
  static uint8_t maskedDataBuf[4096];
  if (length > (int)sizeof(maskedDataBuf)) {
    wsConnected = 0;
    Serial.println("frame too large");
    return;
  }
  for (int i = 0; i < length; i++) {
    maskedDataBuf[i] = buff[i] ^ maskKey[i % 4];
  }
  size_t written = wsClient.write(maskedDataBuf, length);
  
  if (written != length) {
    Serial.printf("send data failed, written: %d, expected: %d\n", written, length);
    wsConnected = 0;
    return;
  }
}